
"""
与数据库建立连接，
自测：运行失败，网络不通
 ('Connection aborted.', ConnectionResetError(54, 'Connection reset by peer'))
"""
from langchain.agents import AgentType
from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_ollama import OllamaLLM

# llm = ChatOpenAI(temperature=0)
llm = OllamaLLM(model="llama3.2:3b")
# load_dotenv()
#
# llm = AzureChatOpenAI(
#     # openai_api_key=
#     # openai_api_base=os.getenv("AZURE_OPENAI_ENDPOINT"),
#     api_key=os.getenv("AZURE_OPENAI_API_KEY"),
#     azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
#     azure_deployment=os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME"),
#     api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
#     temperature=0.7
# )

# db = SQLDatabase.from_uri("sqlite:///local_cache.db")
db = SQLDatabase.from_uri("sqlite:///example.db")
# 大模型做总结，归纳
toolkit = SQLDatabaseToolkit(db=db,llm=llm)
print(toolkit.get_tools())

agent_executor = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    verbose=False,
    agent_executor=AgentType.OPENAI_FUNCTIONS ## 这里要看下llama的fuction 在哪
)

if __name__ == '__main__':
    result = agent_executor.invoke("Describe the full_llm_cache table")
    print(result)